# 1 导入库
import os

os.environ["CUDA_VISIBLE_DEVICES"] = "0"
import torch
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
from torchvision import models
from torch.autograd import Variable

# from net import vgg16_bn
from resnet_yolo import resnet50
from yoloLoss import yoloLoss
from data_label import yoloDataset
from visualizer import Visualizer
import numpy as np

# 2 设置参数
# use_gpu = torch.cuda.is_available()
file_root = r'/Users/liushuang/PycharmProjects/YOLOV1_LS'
learning_rate = 0.001
num_epochs = 2
batch_size = 2
use_resnet = True

# 3 backbone
if use_resnet:
    net = resnet50()
else:
    pass
    # net = vgg16_bn()
print(net)

# 3.1 导入预训练参数
if use_resnet:
    resnet = models.resnet50(pretrained=False)  # True
    new_state_dict = resnet.state_dict()
    dd = net.state_dict()
    for k in new_state_dict.keys():
        if k in dd.keys() and not k.startswith('fc'):
            dd[k] = new_state_dict[k]
    net.load_state_dict(dd)
else:
    vgg = models.vgg16_bn(pretrained=False)
    new_state_dict = vgg.state_dict()
    dd = net.state_dict()
    for k in new_state_dict.keys():
        print(k)
        if k in dd.keys() and k.startswith('features'):
            print('yes')
            dd[k] = new_state_dict[k]
    net.load_state_dict(dd)
if False:
    net.load_state_dict(torch.load('best.pth'))

# 4 Loss
criterion = yoloLoss(7, 2, 5, 0.5)

# if use_gpu:
#     net.cuda()
# 模型训练
net.train()

# 5 参数
params = []
params_dict = dict(net.named_parameters())

for k, v in params_dict.items():
    if k.startswith('features'):
        params += [{'params': [v], 'lr': learning_rate * 1}]
    else:
        params += [{'params': [v], 'lr': learning_rate * 1}]

# 6 优化器
optimizer = torch.optim.SGD(params, lr=learning_rate, momentum=0.9, weight_decay=5e-4)

# 7 导入数据

train_dataset = yoloDataset(root=file_root, list_file='2007_LS_train.txt', train=True,
                            transform=[transforms.ToTensor()])
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=0)

test_dataset = yoloDataset(root=file_root, list_file='2007_LS_test.txt', train=False, transform=[transforms.ToTensor()])
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=4)

print('the dataset has %d images' % (len(train_dataset)))
print('the batch_size is %d' % (batch_size))
logfile = open('log.txt', 'w')

num_iter = 0
# vis = Visualizer()
best_test_loss = np.inf

# 8 训练
for epoch in range(num_epochs):
    net.train()
    if epoch == 30:
        learning_rate = 0.0001
    if epoch == 40:
        learning_rate = 0.00001
    for params_group in optimizer.param_groups:
        params_group['lr'] = learning_rate

    print('\n\nStarting epoch %d / %d' % (epoch + 1, num_epochs))
    print('Learning Rate for this epoch: {}'.format(learning_rate))

    total_loss = 0.

    for i, (images, target) in enumerate(train_loader):
        images = Variable(images)
        target = Variable(target)
        # if use_gpu:
        #     images,target = images.cuda(),target.cuda()

        pred = net(images)
        print('pred.shape:', pred.shape)
        print('target.shape:', target.shape)
        loss = criterion(pred, target)
        total_loss += loss.data.item()

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if (i + 1) % 5 == 0:
            print('Epoch [%d/%d], Iter [%d/%d] Loss: %.4f, average_loss: %.4f'
                  % (epoch + 1, num_epochs, i + 1, len(train_loader), loss.data.item(), total_loss / (i + 1)))
            num_iter += 1
            # vis.plot_train_val(loss_train=total_loss / (i + 1))

    validation_loss = 0.0
    net.eval()
    for i, (images, target) in enumerate(test_loader):
        images = Variable(images, volatile=True)
        target = Variable(target, volatile=True)
        if use_gpu:
            images, target = images.cuda(), target.cuda()
        pred = net(images)
        loss = criterion(pred, target)
        validation_loss += loss.data[0]
    validation_loss /= len(test_loader)
    # vis.plot_train_val(loss_val=validation_loss)

    if best_test_loss > validation_loss:
        best_test_loss = validation_loss
        print('get best test loss %.5f' % best_test_loss)
        torch.save(net.state_dict(), 'best.pth')
    logfile.writelines(str(epoch) + '\t' + str(validation_loss) + '\n')
    logfile.flush()
    torch.save(net.state_dict(), 'yolo.pth')
